import os
import sys
import logging
import paddle
import cv2
import numpy as np
import matplotlib.pylab as plt

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

from ppocr.modeling.architectures import build_model
from ppocr.optimizer import build_optimizer
from ppocr.utils.save_load import init_model
from ppocr.data import SimpleDataSet, build_dataloader
from ppocr.utils.logging import get_logger
from tools import program


def denormalize(img):
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])

    img = np.transpose(img, (1, 2, 0))
    img = img[:, :, ::-1]
    img = img * std + mean
    img = np.clip(img, 0, 1)    

    return img


def show_east():
    cfg_file = "F:/Projects/cv/paddleocr/exercise/det_r50_fast.yml"
    config = program.load_config(cfg_file)

    logger = get_logger()
    dataset = SimpleDataSet(config, 'Train', logger)
    for i in range(10):
        img, score_map, geo_map, mask = dataset[i]
        img = denormalize(img)
        plt.subplot(121)
        plt.imshow(img)
        plt.subplot(122)
        plt.imshow(geo_map[0])
        plt.show()

    # fig, axs = plt.subplots(2, 2, tight_layout={'pad': 1})
    # axs[0, 0].imshow(img)
    # axs[0, 0].imshow(score_map[0], 'hot', alpha=0.2)
    # axs[0, 1].imshow(score_map[0])
    # axs[1, 0].imshow(geo_map[3])
    # axs[1, 1].imshow(mask[0])
    # for i in range(9):
    #     plt.subplot(121)
    #     plt.title(str(i))
    #     plt.imshow(img)
    #     plt.subplot(122)
    #     plt.imshow(geo_map[i])
    #     plt.show()



if __name__ == '__main__':
    show_east()